import queue
from app.model.YoloModel import YoloModel


class YoloService:
    """
    A service wrapper for the YoloModel, responsible for managing video capture,
    model inference, and returning various frame outputs to upstream consumers.
    """

    def __init__(self, model_path, src_points, cap, hot_zone=None, stay_threshold=5, traffic_flow=False, num_lanes=2):
        """
        Initialize the YoloService.

        Args:
            model_path (str): Path to the YOLO model file.
            src_points (np.array): Source points for perspective transform.
            cap (cv2.VideoCapture): OpenCV video capture object.
            hot_zone (ndarray, optional): Polygon coordinates for a special zone (e.g., no-parking zone).
            stay_threshold (int): Time (in seconds) after which a vehicle is considered as staying too long in hot zone.
            traffic_flow (bool): Whether to enable traffic flow counting (crossing middle line).
            num_lanes (int): Number of lanes to draw in the bird’s-eye view.
        """
        self.model = YoloModel(model_path, src_points, hot_zone, stay_threshold, traffic_flow, num_lanes)
        self.rowQueue = queue.Queue(maxsize=2)  # Stores original frames
        self.processedQueue = queue.Queue(maxsize=2)  # Stores annotated frames with tracking results
        self.birdViewQueue = queue.Queue(maxsize=2)  # Stores bird's-eye view output
        self.cap = cap

    @staticmethod
    def try_put(q, item):
        """
        Try to put an item into a queue without blocking.

        Args:
            q (queue.Queue): The queue to put the item into.
            item (any): The item to put into the queue.
        """
        try:
            q.put_nowait(item)
        except queue.Full:
            q.get_nowait()  # Discard oldest
            q.put_nowait(item)

    def start(self):
        """
        Start reading frames from video stream and processing them with the model.
        Processed frames are placed in their respective queues.
        """
        try:
            while True:
                ret, frame = self.cap.read()
                if ret:
                    processed, row, birdView = self.model.track(frame)

                    # Queue the results
                    self.try_put(self.processedQueue, processed)
                    self.try_put(self.rowQueue, row)
                    self.try_put(self.birdViewQueue, birdView)
                else:
                    break
        except Exception as e:
            print(f"[YoloService] Error during processing: {e}")
        finally:
            self.release()

    def get_statistics(self):
        """
        Retrieve the current vehicle tracking statistics.

        Returns:
            dict: Includes total count, category breakdown, long stays, and crossings.
        """
        return self.model.get_statistics()

    def get_row_frame(self):
        """
        Get the latest original (unprocessed) frame from the queue.

        Returns:
            np.ndarray or None: The raw frame if available, otherwise None.
        """
        if not self.rowQueue.empty():
            return self.rowQueue.get()
        return None

    def get_processed_frame(self):
        """
        Get the latest annotated frame (with detection and tracking info).

        Returns:
            np.ndarray or None: The processed frame if available, otherwise None.
        """
        if not self.processedQueue.empty():
            return self.processedQueue.get()
        return None

    def get_birdView_frame(self):
        """
        Get the latest bird's-eye view frame showing vehicle trajectories.

        Returns:
            np.ndarray or None: The bird's-eye view frame if available, otherwise None.
        """
        if not self.birdViewQueue.empty():
            return self.birdViewQueue.get()
        return None

    def release(self):
        """
        Release resources including the video capture and reset tracking statistics.
        """
        self.cap.release()
        self.model.reset_statistics()
